package com.dragon.aipictures.test2;

import com.dragon.aipictures.domain.entity.Picture;
import lombok.extern.slf4j.Slf4j;
import org.springframework.jdbc.core.JdbcTemplate;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.*;

@Slf4j
public class DataMigrator {
    private static final int BATCH_SIZE = 2000;
    private final Long spaceId;
    private final DataSource dataSource;

    private JdbcTemplate jdbcTemplate;
    private final int shardCount; // 新增分片数参数
    public DataMigrator(Long spaceId,
                        DataSource dataSource,
                        JdbcTemplate jdbcTemplate,
                        int shardCount) {
        this.spaceId = spaceId;
        this.dataSource = dataSource;
        this.jdbcTemplate = jdbcTemplate;
        this.shardCount = shardCount;
    }
    public void migrate() {
        long total = getTotalCount();
        int threadCount = calculateThreadCount(total);

        ExecutorService executor = Executors.newFixedThreadPool(threadCount);
        List<Future<?>> futures = new ArrayList<>();

        for (int i = 0; i < threadCount; i++) {
            final int threadNo = i;
            futures.add(executor.submit(() -> processBatch(threadNo, threadCount)));
        }

        awaitCompletion(futures);
        executor.shutdown();
    }

    private int calculateThreadCount(long total) {
        // 策略：每个线程处理不超过10万条数据，但不超过CPU核心数*2
        int maxThreads = Runtime.getRuntime().availableProcessors() * 2;
        int basedOnData = (int) Math.ceil((double) total / 100_000);

        return Math.min(maxThreads, Math.max(1, basedOnData));
    }

    // 带动态调整的版本
    private int calculateDynamicThreadCount(long total) {
        int baseThreads = Runtime.getRuntime().availableProcessors();
        long maxMemory = Runtime.getRuntime().maxMemory();

        // 内存保护：每条记录预估1KB
        long estimatedMemory = total * 1024;
        if (estimatedMemory > maxMemory * 0.3) {
            return Math.max(1, baseThreads / 2);
        }

        return Math.min(baseThreads * 4, (int) (total / 10_000) + 1);
    }

    private void processBatch(int threadNo, int totalThreads) {
        try (Connection conn = dataSource.getConnection();
             PreparedStatement selectStmt = conn.prepareStatement(
                 "SELECT id, content, created_at FROM picture_" + spaceId 
                 + " WHERE MOD(id, ?) = ? ORDER BY id")) {
            
            selectStmt.setInt(1, totalThreads);
            selectStmt.setInt(2, threadNo);
            
            ResultSet rs = selectStmt.executeQuery();
            
            List<Picture> buffer = new ArrayList<>(BATCH_SIZE);
            while (rs.next()) {
                buffer.add(mapRow(rs));
                if (buffer.size() >= BATCH_SIZE) {
                    insertBatch(buffer);
                    buffer.clear();
                }
            }
            insertBatch(buffer); // 处理剩余记录
        } catch (SQLException e) {
            throw new RuntimeException("数据迁移失败", e);
        }
    }

    private Picture mapRow(ResultSet rs) throws SQLException {
        Picture picture = new Picture();
        picture.setId(rs.getLong("id"));
        picture.setSpaceId(spaceId);
        return picture;
    }

    private void insertBatch(List<Picture> pictures) {
        if (pictures.isEmpty()) return;
        
        pictures.parallelStream().forEach(pic -> {
            int shard = (int) (pic.getId() % 4);
            String sql = String.format(
                "INSERT INTO picture_%d_%d (id, space_id, content, created_at) "
                + "VALUES (?, ?, ?, ?)",
                spaceId, shard
            );
            
            jdbcTemplate.update(sql, 
                pic.getId(),
                pic.getSpaceId()
            );
        });
    }
    private void awaitCompletion(List<Future<?>> futures) {
        for (Future<?> future : futures) {
            try {
                future.get(30, TimeUnit.MINUTES); // 设置合理超时时间
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();

//                throw new MigrationException("迁移被中断", e);
            } catch (TimeoutException e) {
//                throw new MigrationException("迁移任务超时", e);
            } catch (ExecutionException e) {
//                throw new MigrationException("迁移执行失败", e.getCause());
            }
        }
    }

    // 增强版（带进度跟踪）
    private void awaitCompletionWithProgress(List<Future<?>> futures, long total) {
        long start = System.currentTimeMillis();
        long processed = 0;

        for (Future<?> future : futures) {
            try {
                future.get();
                processed += BATCH_SIZE;

                // 打印进度
                if (System.currentTimeMillis() - start > 5000) { // 每5秒打印一次
                    log.info("迁移进度: {}/{} ({}%)",
                            processed, total,
                            String.format("%.1f", processed * 100.0 / total));
                    start = System.currentTimeMillis();
                }
            } catch (Exception e) {
                // 异常处理...
            }
        }
    }
    private long getTotalCount() {
        String sourceTable = "picture_" + spaceId;
        return jdbcTemplate.queryForObject(
                "SELECT COUNT(*) FROM " + sourceTable,
                Long.class
        );
    }

    // 带重试机制的版本（推荐生产使用）
//    private long getTotalCountWithRetry() {
//        return RetryTemplate.builder()
//                .maxAttempts(3)
//                .fixedBackoff(1000)
//                .retryOn(DataAccessException.class)
//                .build()
//                .execute(context -> {
//                    return jdbcTemplate.queryForObject(
//                            "SELECT /*+ MAX_EXECUTION_TIME(5000) */ COUNT(*) FROM picture_" + spaceId,
//                            Long.class
//                    );
//                });
//    }
}